/*
*  Copyright (c) 2005-2010, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
*
*  WSO2 Inc. licenses this file to you under the Apache License,
*  Version 2.0 (the "License"); you may not use this file except
*  in compliance with the License.
*  You may obtain a copy of the License at
*
*    http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied.  See the License for the
* specific language governing permissions and limitations
* under the License.
*/

package com.jisc.esb.tasks;

import org.apache.axiom.om.OMElement;
import org.apache.axiom.om.util.AXIOMUtil;
import org.apache.axis2.AxisFault;
import org.apache.axis2.context.ConfigurationContext;
import org.apache.axis2.context.ConfigurationContextFactory;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.synapse.ManagedLifecycle;
import org.apache.synapse.MessageContext;
import org.apache.synapse.SynapseException;
import org.apache.synapse.config.Entry;
import org.apache.synapse.config.SynapseConfiguration;
import org.apache.synapse.core.SynapseEnvironment;
import org.apache.synapse.mediators.builtin.CalloutMediator;
import org.apache.synapse.task.Task;
import org.apache.synapse.util.PayloadHelper;

import javax.xml.namespace.QName;
import javax.xml.stream.XMLStreamException;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

public class HealthChecker implements Task, ManagedLifecycle {

    private static final Log log = LogFactory.getLog(HealthChecker.class);

    private static final QName ENDPOINTS_ELT = new QName("endpoints");
    private static final QName ENDPOINT_ELT = new QName("endpoint");

    private String endpointListKey;
    private String clientRepository = null;
    private String axis2xml = null;
    private String action;
    private OMElement message;

    private OMElement endpointListElt;

    private SynapseEnvironment synapseEnvironment;
    private ConfigurationContext configCtx;
    private ExecutorService workerPool;

    private String faultKey = "";
    private String successKey = "";


    public void execute() {
        SynapseConfiguration synapseConfig = synapseEnvironment.getSynapseConfiguration();
        Entry entry = synapseConfig.getEntryDefinition(endpointListKey);

        // Retrieve the endpoint list from registry
        boolean needsReload = (entry != null) && entry.isDynamic() &&
                              (!entry.isCached() || entry.isExpired());
        if (endpointListElt == null || needsReload) {
            Object endpointList = synapseConfig.getEntry(endpointListKey);
            if (endpointList != null && endpointList instanceof OMElement) {
                endpointListElt = (OMElement) endpointList;
            } else {
                log.warn("Unable to load endpoints from: " + endpointListKey);
                return;
            }
        }

        if (!endpointListElt.getQName().equals(ENDPOINTS_ELT)) {
            log.warn("Invalid parent element: " + endpointListElt.getQName());
            return;
        }

        Iterator endpoints = endpointListElt.getChildrenWithName(ENDPOINT_ELT);
        List<Future<HealthCheckWorker>> futures = new ArrayList<Future<HealthCheckWorker>>();
        while (endpoints.hasNext()) {
            // Submit each heart beat to the worker pool
            OMElement endpointElt = (OMElement) endpoints.next();
            HealthCheckWorker worker = new HealthCheckWorker(endpointElt.getText(), action,
                                                             message.cloneOMElement(), configCtx);

            // Keep track of the futures so we can figure out the status of each heartbeat
            Future<HealthCheckWorker> future = workerPool.submit(worker, worker);
            futures.add(future);
        }

        for (Future<HealthCheckWorker> future : futures) {
            checkStatus(future);
        }
    }

    private void checkStatus(Future<HealthCheckWorker> future) {
        try {
            HealthCheckWorker worker = future.get();
            if (worker.isSuccess()) {
                handleEndpointResponse(worker.getEndpoint(), worker.getResponse());
            } else {
                handleEndpointFault(worker.getEndpoint(), worker.getException());
            }
        } catch (Exception e) {
            log.error("Error while sending a heartbeat", e);
        }
    }

    private void handleEndpointResponse(String endpoint, OMElement response) {
        // TODO: Success case
        log.info("Received response from endpoint: " + endpoint);
        if (!"".equals(successKey)) {
            log.debug("success case for Heartbeat Checker Endpoint: " + endpoint + "  key: " + successKey);
            MessageContext messageCtxt = synapseEnvironment.createMessageContext();
            PayloadHelper.setXMLPayload(messageCtxt, response);
            messageCtxt.setProperty("HB_STATE", successKey);
            messageCtxt.setProperty("HB_ENDPOINT", endpoint);
            synapseEnvironment.injectMessage(messageCtxt);
        } else {
            log.debug("No success case defined for Heartbeat Checker: ");
        }
    }

    private void handleEndpointFault(String endpoint, Exception e) {
        // TODO: Fault case
        log.error("Error connecting to endpoint: " + endpoint);
        if (!"".equals(faultKey)) {
            log.debug("success case for Heartbeat Checker Endpoint: " + endpoint + "  key: " + faultKey);
            MessageContext messageCtxt = synapseEnvironment.createMessageContext();
            OMElement errorMsg = createErrorResponse(e);
            PayloadHelper.setXMLPayload(messageCtxt, errorMsg);
            messageCtxt.setProperty("HB_STATE", faultKey);
            messageCtxt.setProperty("HB_ENDPOINT", endpoint);
            messageCtxt.setProperty("HB_ERROR", e.getMessage());
            synapseEnvironment.injectMessage(messageCtxt);
        } else {
            log.debug("No failure case defined for Heartbeat Checker: ");
        }

    }

    private OMElement createErrorResponse(Exception e) {
        String errorString = "<HeartBeatError><message>" + e.getMessage() + "</message></HeartBeatError>";
        OMElement errorResponse = null;
        try {
            errorResponse = AXIOMUtil.stringToOM(errorString);
        } catch (XMLStreamException e1) {
            return  null;
        }
        return errorResponse;
    }

    public void init(SynapseEnvironment synapseEnvironment) {
        if (endpointListKey == null) {
            error("Endpoint list key not specified", null);
        } else if (message == null) {
            error("Message payload not specified", null);
        }

        if (log.isDebugEnabled()) {
            log.debug("Initializing health checker task");
        }

        this.synapseEnvironment = synapseEnvironment;
        try {
            this.configCtx = ConfigurationContextFactory.createConfigurationContextFromFileSystem(
                    clientRepository != null ? clientRepository : CalloutMediator.DEFAULT_CLIENT_REPO,
                    axis2xml != null ? axis2xml : CalloutMediator.DEFAULT_AXIS2_XML);
        } catch (AxisFault axisFault) {
            error("Error while initializing client configuration context", axisFault);
        }

        workerPool = Executors.newCachedThreadPool(new HealthCheckerThreadFactory());
    }

    public void destroy() {
        if (log.isDebugEnabled()) {
            log.debug("Cleaning up health checker task");
        }

        synapseEnvironment = null;
        try {
            configCtx.terminate();
        } catch (AxisFault axisFault) {
            log.warn("Error while terminating configuration context", axisFault);
        } finally {
            configCtx = null;
        }
        workerPool.shutdownNow();
    }

    private void error(String msg, Exception e) {
        if (e != null) {
            log.error(msg, e);
            throw new SynapseException(msg, e);
        } else {
            log.error(msg);
            throw new SynapseException(msg);
        }
    }

    public String getEndpointListKey() {
        return endpointListKey;
    }

    public void setEndpointListKey(String endpointListKey) {
        this.endpointListKey = endpointListKey;
    }

    public String getClientRepository() {
        return clientRepository;
    }

    public void setClientRepository(String clientRepository) {
        log.info("Using custom Axis2 repository path: " + clientRepository);
        this.clientRepository = clientRepository;
    }

    public String getAxis2xml() {
        return axis2xml;
    }

    public void setAxis2xml(String axis2xml) {
        log.info("Using custom axis2.xml file: " + axis2xml);
        this.axis2xml = axis2xml;
    }

    public String getAction() {
        return action;
    }

    public void setAction(String action) {
        this.action = action;
    }

    public OMElement getMessage() {
        return message;
    }

    public void setMessage(OMElement message) {
        this.message = message;
    }

    public String getFaultKey() {
        return faultKey;
    }

    public void setFaultKey(String faultKey) {
        this.faultKey = faultKey;
    }

    public String getSuccessKey() {
        return successKey;
    }

    public void setSuccessKey(String successKey) {
        this.successKey = successKey;
    }


}
